'''
    保温杯基础类
'''
import pdb
import dataclasses
import importlib
from typing import Dict, List, Tuple, Union

import pandas as pd
from joblib import Parallel, delayed

from alphalib.backtest.playback import timer
from alphalib.common_utils import f1_utils
from alphalib.config_backtest import factor_path, filter_path
from alphalib.contrib.strategy.base_strategy import Z1BaseStrategy
from alphalib.contrib.strategy.z1_conf_types import CalcFactorType, F1FactorConfig, F1FilterParams, factor_to_col_name


@dataclasses.dataclass
class Z1StrategyConfig:
    strategy_name: str = 'STRATEGY_NAME'
    hold_period: int = 12
    calc_factor_type: CalcFactorType = 'cross'
    long_factors: List[F1FactorConfig] = dataclasses.field(default_factory=list)
    short_factors: List[F1FactorConfig] = dataclasses.field(default_factory=list)
    filter_before_params: List[F1FilterParams] = dataclasses.field(default_factory=list)
    filter_after_params: List[F1FilterParams] = dataclasses.field(default_factory=list)
    if_use_spot: bool = True
    long_weight: float = 1
    short_weight: float = 1
    long_coin_num: Union[int, float] = 1
    short_coin_num: Union[int, float] = 1
    stg_weight: float = 1
    enable_seal999: bool = False

    long_white_list: List = dataclasses.field(default_factory=list)
    short_white_list: List = dataclasses.field(default_factory=list)
    long_black_list: List = dataclasses.field(default_factory=list)
    short_black_list: List = dataclasses.field(default_factory=list)


class Z1Strategy(Z1BaseStrategy):
    z1_strategy_config: Z1StrategyConfig

    def __init__(self, z1_strategy_config: Z1StrategyConfig):
        self.z1_strategy_config = z1_strategy_config
        if abs(self.z1_strategy_config.long_weight + self.z1_strategy_config.short_weight - 2) > 1e-8:
            print(
                '建议设置`long_weight+short_weight==2`。'
                '可通过`F1PlaybackConfig.leverage`控制杠杆，通过`F1StrategyConfig.stg_weight控制组合权重`'
            )
            print(f'当前 long_weight={z1_strategy_config.long_weight} short_weight={z1_strategy_config.short_weight}')

    def calc_alloc_ratio(self, all_spot_df: Dict[str, pd.DataFrame], all_swap_df: Dict[str, pd.DataFrame], start_date: str, end_date: str, n_jobs: int = 1) -> pd.Series:
        # columns = ["symbol", "candle_begin_time", "open", "high", "low", "close", "volume", "quote_volume",
        #            "trade_num", "taker_buy_base_asset_volume", "taker_buy_quote_asset_volume", "fundingRate", "funding_rate_raw"]
        # columns = ["symbol", "candle_begin_time", "open", "high", "low", "close", "volume", "quote_volume",
        #            "trade_num", "taker_buy_base_asset_volume", "taker_buy_quote_asset_volume"]
        start_date = pd.to_datetime(start_date)
        end_date = pd.to_datetime(end_date)
        all_spot_df = {symbol: df.copy(deep=False) for symbol, df in all_spot_df.items()}
        all_swap_df = {symbol: df.copy(deep=False) for symbol, df in all_swap_df.items()}
        long_factor_conf_list = list(set(self.z1_strategy_config.long_factors))
        short_factor_conf_list = list(set(self.z1_strategy_config.short_factors))
        long_filter_conf_list = list(set(x.to_filter_config()
            for x in self.z1_strategy_config.filter_before_params + self.z1_strategy_config.filter_after_params
        ))
        short_filter_conf_list = list(set(x.to_filter_config()
            for x in self.z1_strategy_config.filter_before_params + self.z1_strategy_config.filter_after_params
        ))

        if self.z1_strategy_config.if_use_spot:
            all_long_df = all_spot_df
            all_short_df = all_swap_df
        else:
            all_long_df = all_swap_df
            all_short_df = all_swap_df

        with timer.timer('--calc factors'):
            long_df_factor_merge, long_trade_mask = self.cal_one_side(all_long_df, long_factor_conf_list, long_filter_conf_list, start_date, end_date)
            short_df_factor_merge, short_trade_mask = self.cal_one_side(all_short_df, short_factor_conf_list, short_filter_conf_list, start_date, end_date)

        # 前置过滤
        long_inds, short_inds = self.filter_handle_before(long_df_factor_merge[long_trade_mask], short_df_factor_merge[short_trade_mask])

        if self.z1_strategy_config.long_factors == self.z1_strategy_config.short_factors:
            calc_factor_list = [self.z1_strategy_config.long_factors]
            col_list = ['因子', '因子']
        else:
            calc_factor_list = [self.z1_strategy_config.long_factors, self.z1_strategy_config.short_factors]
            col_list = ['多头因子', '空头因子']

        # 多空一致，数组长度为1，多空分离，数组长度为2
        for factor_conf_list, col in zip(calc_factor_list, col_list):
            factor_list = [factor_to_col_name(x) for x in factor_conf_list]
            factor_coef = pd.Series([(-1 if x.descending else 1) * x.weight for x in factor_conf_list],
                                    index=factor_list)
            if self.z1_strategy_config.calc_factor_type == 'cross':
                if col == '因子' or col == '多头因子':
                    long_df_factor_merge[factor_list] = long_df_factor_merge.groupby('candle_begin_time')[factor_list].rank(
                        ascending=True)
                    long_factor_meds = long_df_factor_merge.groupby('candle_begin_time')[factor_list].transform('median')
                    long_df_factor_merge[factor_list] = long_df_factor_merge[factor_list].fillna(long_factor_meds)
                if col == '因子' or col == '空头因子':
                    short_df_factor_merge[factor_list] = short_df_factor_merge.groupby('candle_begin_time')[factor_list].rank(
                        ascending=True)
                    short_factor_meds = short_df_factor_merge.groupby('candle_begin_time')[factor_list].transform('median')
                    short_df_factor_merge[factor_list] = short_df_factor_merge[factor_list].fillna(short_factor_meds)
            if col == '因子' or col == '多头因子':
                long_df_factor_merge[col] = long_df_factor_merge[factor_list].dot(factor_coef.T)  # use np.dot instead?
            if col == '因子' or col == '空头因子':
                short_df_factor_merge[col] = short_df_factor_merge[factor_list].dot(factor_coef.T)  # use np.dot instead?

        # 计算选币因子
        cols_for_gen_select_long = ['candle_begin_time', 'symbol'] + [col_list[0]]
        cols_for_gen_select_short = ['candle_begin_time', 'symbol'] + [col_list[1]]
        df_long = long_df_factor_merge.loc[long_inds, cols_for_gen_select_long]
        df_short = short_df_factor_merge.loc[short_inds, cols_for_gen_select_short]
        rank_long = df_long.groupby('candle_begin_time')[col_list[0]].rank(
            method='first', pct=self.z1_strategy_config.long_coin_num < 1)
        rank_short = df_short.groupby('candle_begin_time')[col_list[1]].rank(
            method='first', pct=self.z1_strategy_config.short_coin_num < 1, ascending=False)
        select_long_mask = (rank_long <= self.z1_strategy_config.long_coin_num)
        select_short_mask = (rank_short <= self.z1_strategy_config.short_coin_num)
        select_long_inds = df_long.index[select_long_mask]
        select_short_inds = df_short.index[select_short_mask]

        # 后置过滤
        select_long_inds, select_short_inds = self.filter_handle_after(
            long_df_factor_merge.loc[select_long_inds, :],
            short_df_factor_merge.loc[select_short_inds, :]
        )

        # 选币
        if self.z1_strategy_config.long_coin_num < 1:
            long_coin_num = select_long_mask.groupby(df_long['candle_begin_time']).transform(sum)
            long_coin_num = long_coin_num.loc[select_long_inds]
        else:
            long_coin_num = self.z1_strategy_config.long_coin_num
        if self.z1_strategy_config.short_coin_num < 1:
            short_coin_num = select_short_mask.groupby(df_short['candle_begin_time']).transform(sum)
            short_coin_num = short_coin_num.loc[select_short_inds]
        else:
            short_coin_num = self.z1_strategy_config.short_coin_num

        # 计算单offset资金分配比例
        long_base_coef = 0.5 * self.z1_strategy_config.long_weight / self.z1_strategy_config.hold_period
        short_base_coef = 0.5 * self.z1_strategy_config.short_weight / self.z1_strategy_config.hold_period

        long_df_factor_merge['alloc_ratio_1offset'] = 0
        short_df_factor_merge['alloc_ratio_1offset'] = 0
        long_df_factor_merge.loc[select_long_inds, 'alloc_ratio_1offset'] += long_base_coef / long_coin_num
        short_df_factor_merge.loc[select_short_inds, 'alloc_ratio_1offset'] -= short_base_coef / short_coin_num

        # 按照hold_period合并coef
        # 在f1框架的选币逻辑中，给每个offset分配资金后，是按照那个offset当时的close价格计算 target_amount = assign_usdt / close，
        # 相当于乘上调整系数 close_最新 / close_当时
        with timer.timer('--merge alloc_ratio'):
            long_df_factor_merge['alloc_ratio_1offset'] /= long_df_factor_merge['close']
            long_df_factor_merge['alloc_ratio_long'] = long_df_factor_merge.groupby('symbol', group_keys=False)[
                'alloc_ratio_1offset'].rolling(
                self.z1_strategy_config.hold_period, min_periods=1).sum().reset_index(level=0, drop=True)
            long_df_factor_merge['alloc_ratio_long'] *= long_df_factor_merge['close']
            short_df_factor_merge['alloc_ratio_1offset'] /= short_df_factor_merge['close']
            short_df_factor_merge['alloc_ratio_short'] = short_df_factor_merge.groupby('symbol', group_keys=False)[
                'alloc_ratio_1offset'].rolling(
                self.z1_strategy_config.hold_period, min_periods=1).sum().reset_index(level=0, drop=True)
            short_df_factor_merge['alloc_ratio_short'] *= short_df_factor_merge['close']

        long_res = long_df_factor_merge[['candle_begin_time', 'symbol', 'alloc_ratio_long']]
        short_res = short_df_factor_merge[['candle_begin_time', 'symbol', 'alloc_ratio_short']]
        res = pd.merge(long_res, short_res, on=['candle_begin_time', 'symbol'], how='outer').fillna(0)
        res['alloc_ratio'] = res['alloc_ratio_long'] + res['alloc_ratio_short']
        res = res.set_index(['candle_begin_time', 'symbol'])
        res = res['alloc_ratio']
        cbt_index = res.index.levels[0]
        cbt_index = cbt_index.sort_values()
        start_date_idx = cbt_index.searchsorted(pd.to_datetime(start_date))
        start_date_idx = max(start_date_idx - 1, 0)
        res_start_date = cbt_index[start_date_idx]
        res = res.loc[res.index.get_level_values('candle_begin_time') >= res_start_date]
        return res

    def cal_one_side(self, all_df: Dict[str, pd.DataFrame], factor_conf_list: List, filter_conf_list: List,
                     start_date: pd.Series, end_date: pd.Series) -> Tuple[pd.DataFrame, pd.Series]:
        all_df_factor = {}
        for symbol, df in all_df.items():
            df.loc[df['swap_symbol'] != "", 'symbol'] = df['swap_symbol']
            factors_dict = {}
            for factor_conf in factor_conf_list:
                factor_name = factor_conf.factor_name
                back_periods = factor_conf.back_periods
                fdiff_n = factor_conf.fdiff_n
                factor_col_name = factor_to_col_name(factor_conf)
                factor_mod = importlib.import_module(f'{factor_path}.{factor_name}')
                factors_dict[factor_col_name] = factor_mod.signal(df, back_periods, fdiff_n, factor_col_name)[
                    factor_col_name]

            for filter_factor_conf in filter_conf_list:
                factor_name = filter_factor_conf.factor_name
                params = filter_factor_conf.params
                filter_col_name = filter_factor_conf.to_col_name()
                filter_mod = importlib.import_module(f'{filter_path}.{factor_name}')
                factors_dict[filter_col_name] = filter_mod.signal(df, params, filter_col_name)[filter_col_name]

            factors_dict = {
                'candle_begin_time': df['candle_begin_time'],
                'symbol': df['symbol'],
                'close': df['close'],
                **factors_dict,
            }
            df_factor = pd.DataFrame(data=factors_dict)
            if 1:  # 最好在选币时处理vol过滤
                mask = df['volume'] > 0
                if 'avg_price' in df.columns:  # 回测时过滤下周期avg_price为空的数据，实盘时不处理
                    mask_notna = df['avg_price'].shift(-1).notna()
                    mask = mask & mask_notna
            if self.z1_strategy_config.enable_seal999:  # 最好通过前置过滤处理999封印
                df_factor = df_factor.iloc[999:][mask[999:]]
            else:
                df_factor = df_factor[mask]

            start_idx = df_factor['candle_begin_time'].searchsorted(pd.to_datetime(start_date))
            start_idx = max(start_idx - self.z1_strategy_config.hold_period, 0)
            end_idx = df_factor['candle_begin_time'].searchsorted(pd.to_datetime(end_date), side='right')
            df_factor = df_factor.iloc[start_idx:end_idx]

            all_df_factor[symbol] = df_factor

        df_factor_merge = pd.concat(all_df_factor.values(), ignore_index=True, copy=False)
        df_factor_merge = df_factor_merge.sort_values(by=['candle_begin_time', 'symbol'])

        # 提前排除退市币种
        # 首先获取可能退市的币种
        max_time = df_factor_merge['candle_begin_time'].max()
        quit_df = df_factor_merge.groupby('symbol')['candle_begin_time'].max()
        quit_symbols = quit_df[quit_df < max_time].index
        # 退市币种的处理，实盘提前N小时加入黑名单
        quit_time_thresholds = quit_df.loc[quit_symbols] - pd.Timedelta(
            hours=self.z1_strategy_config.hold_period + 1)
        symbol_to_threshold = quit_time_thresholds.to_dict()
        if symbol_to_threshold:
            ths_time = df_factor_merge['symbol'].map(symbol_to_threshold)
            # 使用掩码来筛选数据
            quit_mask = df_factor_merge['symbol'].isin(quit_symbols)
            final_quit_mask = quit_mask & (df_factor_merge['candle_begin_time'] <= ths_time)
            trade_mask = ~quit_mask | final_quit_mask
        else:
            trade_mask = pd.Series(True, index=df_factor_merge.index)

        return (df_factor_merge, trade_mask)


    def filter_handle_before(self, long_df: pd.DataFrame, short_df: pd.DataFrame) -> Tuple[pd.Index, pd.Index]:
        df1_mask = ~long_df['symbol'].isin(self.z1_strategy_config.long_black_list)
        if self.z1_strategy_config.long_white_list:
            df1_mask = long_df['symbol'].isin(self.z1_strategy_config.long_white_list) & df1_mask

        df2_mask = ~short_df['symbol'].isin(self.z1_strategy_config.short_black_list)
        if self.z1_strategy_config.short_white_list:
            df2_mask = short_df['symbol'].isin(self.z1_strategy_config.short_white_list) & df2_mask

        filter_before_params = [
            [x.direction, x.filter_factor, x.filter_type, x.compare_operator, x.filter_value,
             x.rank_ascending, x.filter_after]
            for x in self.z1_strategy_config.filter_before_params
        ]
        filter_before_exec = [f1_utils.filter_generate(param=param) for param in filter_before_params]
        with timer.timer('--do_filter：'):
            df1, df2 = f1_utils.do_filter(long_df[df1_mask], short_df[df2_mask], filter_before_exec)
        return df1.index, df2.index

    def filter_handle_after(self, df1: pd.DataFrame, df2: pd.DataFrame) -> Tuple[pd.Index, pd.Index]:
        filter_after_params = [
            [x.direction, x.filter_factor, x.filter_type, x.compare_operator, x.filter_value,
             x.rank_ascending, x.filter_after]
            for x in self.z1_strategy_config.filter_after_params
        ]
        filter_after_exec = [f1_utils.filter_generate(param=param) for param in filter_after_params]
        df1 = df1.copy(deep=False)
        df2 = df2.copy(deep=False)
        df1['weight_ratio'] = 1
        df2['weight_ratio'] = 1
        df1, df2 = f1_utils.do_filter(df1, df2, filter_after_exec)
        return df1.index[df1['weight_ratio'] == 1], df2.index[df2['weight_ratio'] == 1]

@dataclasses.dataclass
class Z1MultiStrategyConfig:
    multi_strategy_name: str = 'MULTI_STRATEGY_NAME'
    multi_strategy_config: List[Z1StrategyConfig] = dataclasses.field(default_factory=list)


class Z1MultiStrategy(Z1BaseStrategy):
    stg_list: List[Z1Strategy]

    def __init__(self, stg_list: List[Z1Strategy]):
        self.stg_list = stg_list

    def calc_alloc_ratio(self, all_spot_df: Dict[str, pd.DataFrame], all_swap_df: Dict[str, pd.DataFrame], start_date: str, end_date: str, n_jobs: int = 1) -> pd.Series:
        if n_jobs == 1:
            res_list = [
                stg.calc_alloc_ratio(all_spot_df=all_spot_df, all_swap_df=all_swap_df, start_date=start_date, end_date=end_date)
                for stg in self.stg_list
            ]
        else:
            res_list = Parallel(n_jobs=n_jobs)(
                delayed(stg.calc_alloc_ratio)(
                    all_spot_df=all_spot_df, all_swap_df=all_swap_df, start_date=start_date, end_date=end_date
                )
                for stg in self.stg_list
            )
        weight_list = [stg.z1_strategy_config.stg_weight for stg in self.stg_list]
        res = (pd.concat(res_list, axis=1) * weight_list).sum(axis=1) / sum(weight_list)

        return res